import pathlib

import matplotlib.pyplot
import numba
import numpy

from problem1.checkpoint import Checkpoint
from utilities.checkpoint_manager import CheckpointManager


@numba.jit(nopython=True)
def _scan_numba(caloric):
    result = numpy.zeros((1000, 1000))
    for c in caloric:
        x, y = c
        result[x + 500][y + 500] += 1
    return result


def main(checkpoint_manager: CheckpointManager,
         plots: tuple[tuple[int], ...] = ((0, 100, 200), (300, 400, 500))):
    total_row = len(plots)
    total_column = len(plots[0])
    for row_index, row in enumerate(plots):
        for column_index, t in enumerate(row, start=1):
            matplotlib.pyplot.subplot(
                total_row, total_column,
                row_index * total_column + column_index)

            checkpoint = Checkpoint.load(checkpoint_manager.existed()[t])
            checkpoint = _scan_numba(checkpoint.calculation_state())
            checkpoint = checkpoint[440:560, 440:560]
            x = numpy.linspace(-1, 1, checkpoint.shape[0])
            matplotlib.pyplot.contour(
                *numpy.meshgrid(x, x), checkpoint,
                levels=[-1, 5], cmap="rainbow")
            matplotlib.pyplot.contourf(
                *numpy.meshgrid(x, x), checkpoint,
                levels=200, vmin=0, vmax=15000, cmap="hot")
            matplotlib.pyplot.colorbar()
            matplotlib.pyplot.xticks([])
            matplotlib.pyplot.yticks([])
            # noinspection SpellCheckingInspection
            matplotlib.pyplot.title(rf"$t={t} (\mathrm{{Ticks}})$")
    matplotlib.pyplot.subplots_adjust(
        left=0.038, bottom=0.026, right=0.964, top=0.955,
        wspace=0.217, hspace=0.136)
    matplotlib.pyplot.show()


if __name__ == "__main__":
    main(CheckpointManager(pathlib.Path("./outputs")))
